import onnxruntime as ort
import cv2
import numpy as np

def infer_by_onnx(model_path, image_path):
    session = ort.InferenceSession(model_path)

    image = cv2.imread(image_path)
    image_height, image_width, _ = image.shape

    # 假设模型要求的输入是 [1, 3, 384, 640]（需要根据模型实际要求修改）
    input_size = (640, 384)  # 输入大小调整为640x384
    image_resized = cv2.resize(image, input_size)
    image_rgb = cv2.cvtColor(image_resized, cv2.COLOR_BGR2RGB)  # 转换为RGB
    image_normalized = image_rgb.astype(np.float32) / 255.0  # 归一化

    # 需要增加一个维度，以符合模型输入格式 [batch_size, channels, height, width]
    input_tensor = np.expand_dims(image_normalized, axis=0)  # 形状变为 (1, 384, 640, 3)

    # 3. 转换为 [batch_size, channels, height, width]，即 [1, 3, 384, 640]
    input_tensor = np.transpose(input_tensor, (0, 3, 1, 2))  # 形状变为 (1, 3, 384, 640)

    # 4. 推理过程
    # 获取输入的名称（通常是 "input"）
    input_name = session.get_inputs()[0].name
    # 获取输出的名称（在这里我们假设它的名称是 "final_output"）
    output_name = session.get_outputs()[0].name

    # 进行推理
    outputs = session.run([output_name], {input_name: input_tensor})

    # 输出维度为 [1, 1, 10, 25]，解析为：1个样本，1个类别，10个检测框，每个框25个数据项
    output = outputs[0]
    print("Model output shape:", output.shape)
    print("Model output:", output)

    # 5. 解读输出（框信息、类别置信度、类别ID和关键点）
    boxes = output[0, :, :4]  # [10, 4] 每个框的 [center_x, center_y, width, height]
    confidences = output[0, :, 4]  # [10] 每个框的最大类别置信度
    class_ids = output[0, :, 5]  # [10] 每个框的类别 ID
    keypoints = output[0, :, 6:].reshape(10, 6, 3)  # [10, 6, 3] 每个框的 6 个关键点，每个关键点 (x, y, visible)

    # 6. 设置绘图参数
    image_with_boxes = image.copy()

    for i in range(output.shape[1]):  # 10 个框
        box = boxes[i]
        confidence = confidences[i]
        class_id = int(class_ids[i])

        # 计算框的坐标
        x_min = int((box[0] - box[2] / 2) / input_size[0] * image_width)
        y_min = int((box[1] - box[3] / 2) / input_size[1] * image_height)
        x_max = int((box[0] + box[2] / 2) / input_size[0] * image_width)
        y_max = int((box[1] + box[3] / 2) / input_size[1] * image_height)
        print(f"{i} {(x_min, y_min)} {(x_max, y_max)}")

        # 绘制框
        cv2.rectangle(image_with_boxes, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)

        # 绘制置信度和类别
        cv2.putText(image_with_boxes, f"Class {class_id} Conf: {confidence:.2f}",
                    (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # 绘制关键点
        for j in range(6):  # 6 个关键点
            keypoint = keypoints[i, j]
            keypoint_x = int(keypoint[0] / input_size[0] * image_width)
            keypoint_y = int(keypoint[1] / input_size[1] * image_height)

            # 绘制关键点
            cv2.circle(image_with_boxes, (keypoint_x, keypoint_y), 3, (255, 0, 0), -1)
            cv2.putText(image_with_boxes, f"{j + 1}", (keypoint_x + 5, keypoint_y + 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # 7. 显示图像
    cv2.imshow('Detected Boxes and Keypoints', image_with_boxes)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
